import time
import rtde_receive
import rtde_control
from compensation import GravityCompensation
from gym_pih.hardware.KWFT import ftsensor
import threading
import matplotlib.pyplot as plt
import numpy as np

ip = '192.168.3.140'
rtde_r = rtde_receive.RTDEReceiveInterface(ip)
rtde_c = rtde_control.RTDEControlInterface(ip)
ft = ftsensor(False)
compensation = GravityCompensation()

# 定义全局变量，用于保存数据
force_data = []  
compensated_force_data = []
move_done = False

def move():
    target_pos = [0.4883837410367999,
    -0.4299285355317489,
    0.23710715340885336,
    0.002009308761598123,
    -2.8499945615195372,
    0.8602965734054271]  # 请确保这里设置了正确的目标位置
    rtde_c.moveL(target_pos, 0.1, 0.05)
    global move_done
    move_done = True

def record():
    global force_data, compensated_force_data, move_done
    while not move_done:
        force = ft.GetForce()
        rotation_vector = rtde_r.getActualTCPPose()
        # 假设 rotation_vector[3:6] 是所需的旋转信息
        compensated_force = compensation.GetForceTorque(force, rotation_vector[3:6])
        force_data.append(force)
        compensated_force_data.append(compensated_force)
        time.sleep(0.002)  # 根据实际需要调整采样时间

if __name__ == "__main__":
    try:
        rtde_c.moveL([0.5626520164166212,
        -0.09710945663454698,
        0.19297750945266434,
        0.030348783608821865,
        2.963109334006964,
        0.7239321394228534])
        move_thread = threading.Thread(target=move)
        record_thread = threading.Thread(target=record)

        move_thread.start()
        record_thread.start()

        move_thread.join()
        record_thread.join()

        # 数据采集完成后，将数据转换为 numpy 数组
        force_data_np = np.array(force_data)           # 形状: (N, 6)
        comp_force_data_np = np.array(compensated_force_data)  # 形状: (N, 6)
        sample_count = np.arange(force_data_np.shape[0])    # 数据点数量

        # 设置每个维度的 y 轴标签（前三个为力，后三个为扭矩）
        y_labels = ["Fx (N)", "Fy (N)", "Fz (N)", "Mx (Nm)", "My (Nm)", "Mz (Nm)"]

        # 创建3行2列的子图
        fig, axes = plt.subplots(3, 2, figsize=(12, 10))
        axes = axes.flatten()  # 转换为一维数组便于迭代

        for i in range(6):
            ax = axes[i]
            ax.plot(sample_count, force_data_np[:, i], label="Original Force")
            ax.plot(sample_count, comp_force_data_np[:, i], label="Compensated Force")
            ax.set_ylabel(y_labels[i])
            ax.set_xlabel("Data Count")
            ax.legend()
            ax.grid(True)

        fig.tight_layout()
        plt.show()
        
    except KeyboardInterrupt:
        print("Process interrupted by user.")
        rtde_c.stopScript()
